package com.gxc.tree;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集
 */
public class UnionFindSet<V> {

    //对象包装成同一类元素，方便比较
    public HashMap<V, Element<V>> elementMap;
    //key 某个元素， value 该元素的父级
    public HashMap<Element<V>, Element<V>> fatherMap;
    //key 某个元素， value 该元素集合的大小
    public HashMap<Element<V>, Integer> sizeMap;

    public UnionFindSet(List<V> list) {
        elementMap = new HashMap<>();
        fatherMap = new HashMap<>();
        sizeMap = new HashMap<>();

        for (int i = 0; i < list.size(); i++) {
            V v = list.get(i);
            Element<V> element = new Element<V>(v);
            elementMap.put(v, element);
            fatherMap.put(element, element);
            sizeMap.put(element, 1);
        }
    }

    /**
     * 找到元素的最上级的父级（最上级的元素的父级就是自己）
     * 同时修改该元素的所有的父级都指向最上级的父级（降低链路长度）
     * @param element
     * @return
     */
    public Element<V> findHead(Element<V> element) {
        Stack<Element<V>> stack = new Stack<>();
        while (fatherMap.get(element) != element) {
            stack.push(element);
            element = fatherMap.get(element);
        }

        while (!stack.isEmpty()) {
            Element<V> pop = stack.pop();
            fatherMap.put(pop, element);
        }
        return element;
    }

    /**
     * 是否同一个集合
     * @param x
     * @param y
     * @return
     */
    public boolean isSameSet(V x, V y) {
        if (elementMap.containsKey(x) && elementMap.containsKey(y)) {
            return findHead(elementMap.get(x)) == findHead(elementMap.get(y));
        }
        return false;
    }

    private void union(V x, V y) {
        if (elementMap.containsKey(x) && elementMap.containsKey(y)) {
            Element<V> xhead = findHead(elementMap.get(x));
            Element<V> yhead = findHead(elementMap.get(y));
            if (xhead != yhead) {
                Element<V> big = sizeMap.get(xhead) > sizeMap.get(yhead)?xhead:yhead;
                Element<V> small = big==xhead?yhead:xhead;

                fatherMap.put(small, big);
                sizeMap.put(big, sizeMap.get(big) + sizeMap.get(small));
                sizeMap.remove(small);
            }
        }
    }
}
